# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# PyTorch (https://pytorch.org/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
RunOneStepCell
"""

from typing import Tuple, List
import torch
import torch.nn.functional as F
from torch import Tensor
from torch import nn
import logging

from .energy import WithEnergyCell
from .force import WithForceCell
from ...function.functions import get_integer, all_none, get_arguments
from ...optimizer import Updater

class RunOneStepCell(nn.Module):
    r"""
    Module to run one step simulation.
    This Module wraps the `energy` and `force` with the `optimizer`. The backward graph will be created
    in the forward function to update the atomic coordinates of the simulation system.

    Args:
        energy(WithEnergyCell): Module that wraps the simulation system with the potential energy function.
            Defatul: None
        force(WithForceCell): Module that wraps the simulation system with the atomic force function.
            Defatul: None
        optimizer(torch.optim.Optimizer): Optimizer for simulation. Defatul: None
        sink_cycles(int): Cycle steps for sink mode. Default: 1
        sens(float): The scaling number to be filled as the input of backpropagation. Default: 1.0

    Inputs:
        - **\*inputs** (Tuple(Tensor)) - Tuple of input tensors of `WithEnergyCell`.

    Outputs:
        - energy, Tensor of shape `(B, 1)`. Data type is float. Total potential energy.
        - force, Tensor of shape `(B, A, D)`. Data type is float. Atomic force.

    Supported Platforms:
        ``CPU`` ``GPU``

    Symbols:
        B:  Batchsize, i.e. number of walkers of the simulation.
        A:  Number of the atoms in the simulation system.
        D:  Spatial dimension of the simulation system. Usually is 3.
    """
    def __init__(self,
                 energy: WithEnergyCell = None,
                 force: WithForceCell = None,
                 optimizer: torch.optim.Optimizer = None,
                 sens: float = 1.0,
                 **kwargs
                 ):

        super().__init__()
        self._kwargs = get_arguments(locals(), kwargs)

        if all_none([energy, force]):
            raise ValueError('energy and force cannot be both None!')

        self._neighbour_list_pace = None

        self.system_with_energy = energy
        if self.system_with_energy is not None:
            self.system = self.system_with_energy.system
            self.units = self.system_with_energy.units
            self._neighbour_list_pace = self.system_with_energy.neighbour_list_pace

        self.system_with_force = force
        if self.system_with_force is not None:
            force_pace = self.system_with_force.neighbour_list_pace

            if self.system_with_energy is None or self._neighbour_list_pace == 0:
                self._neighbour_list_pace = force_pace
            else:
                if force_pace not in (self._neighbour_list_pace, 0):
                    raise ValueError(
                        f'The neighbour_list_pace in WithForceCell ({force_pace}) cannot match '
                        f'the neighbour_list_pace in WithEnergyCell ({self._neighbour_list_pace}).')

        if self.system_with_energy is None:
            self.system = self.system_with_force.system

        self.optimizer = optimizer
        if self.optimizer is None:
            print('No optimizer! The simulation system will not be updated!')

        self.use_updater = isinstance(self.optimizer, Updater)

    @property
    def neighbour_list_pace(self) -> int:
        r"""
        update step for neighbour list.

        Return:
            int, the number of steps needed for neighbour list updating.
        """
        return self._neighbour_list_pace

    @property
    def energy_cutoff(self) -> Tensor:
        r"""
        cutoff distance for neighbour list in WithEnergyCell.

        Return:
            Tensor, cutoff distance for neighbour list in WithEnergyCell.
        """
        if self.system_with_energy is None:
            return None
        return self.system_with_energy.cutoff

    @property
    def force_cutoff(self) -> Tensor:
        r"""
        cutoff distance for neighbour list in WithForceCell.

        Return:
            Tensor, cutoff distance for neighbour list in WithForceCell.
        """
        if self.system_with_force is None:
            return None
        return self.system_with_force.cutoff

    @property
    def length_unit(self) -> str:
        r"""
        length unit.

        Return:
            str, length unit.
        """
        return self.units.length_unit

    @property
    def energy_unit(self) -> str:
        r"""
        energy unit.

        Return:
            str, energy unit.
        """
        return self.units.energy_unit

    @property
    def num_energies(self) -> int:
        r"""
        number of energy terms :math:`U`.

        Return:
            int, number of energy terms.
        """
        if self.system_with_energy is None:
            return 0
        return self.system_with_energy.num_energies

    @property
    def energy_names(self):
        r"""
        names of energy terms.

        Return:
            list[str], names of energy terms.
        """
        if self.system_with_energy is None:
            return []
        return self.system_with_energy.energy_names

    @property
    def bias_names(self) -> List[str]:
        r"""
        names of bias potential terms.

        Return:
            list[str], names of bias potential terms.
        """
        if self.system_with_energy is None:
            return []
        return self.system_with_energy.bias_names

    @property
    def num_biases(self) -> int:
        r"""
        number of bias potential terms :math:`V`.

        Return:
            int, number of bias potential terms.
        """
        if self.system_with_energy is None:
            return 0
        return self.system_with_energy.num_biases

    @property
    def energies(self) -> Tensor:
        r"""
        energy tensor.

        Return:
            Tensor, energy tensor.
        """
        return self.get_energies()

    @property
    def biases(self) -> Tensor:
        r"""
        biases tensor.

        Return:
            Tensor, biases tensor.
        """
        return self.get_biases()

    @property
    def bias(self) -> Tensor:
        r"""
        bias tensor.

        Return:
            Tensor, bias tensor.
        """
        return self.get_bias()

    @property
    def bias_function(self) -> nn.Module:
        r"""
        bias function.

        Return:
            Cell, bias function.
        """
        if self.system_with_energy is None:
            return None
        return self.system_with_energy.bias_function

    def get_energy(self) -> Tensor:
        r"""
        get energy tensor.
        Return:
            Tensor, energy tensor.
        """
        if self.system_with_force is not None:
            return self.system_with_force.get_energy()
        elif self.system_with_energy is not None:
            return self.system_with_energy.get_energy()
        else:
            return None
    
    def get_energies(self) -> Tensor:
        r"""
        get energy tensor.

        Return:
            Tensor, energy tensor.
        """
        if self.system_with_energy is None:
            return None
        return self.system_with_energy.get_energies()

    def get_biases(self) -> Tensor:
        r"""
        get biases tensor.

        Return:
            Tensor, biases tensor.
        """
        if self.system_with_energy is None:
            return None
        return self.system_with_energy.get_biases()

    def get_bias(self) -> Tensor:
        r"""
        get bias tensor.

        Return:
            Tensor, bias tensor.
        """
        if self.system_with_energy is None:
            return None
        return self.system_with_energy.get_bias()

    def update_neighbour_list(self, steps) -> Tuple[Tensor, Tensor]:
        r"""
        update neighbour list.

        Return:
            tuple[Tensor, Tensor], updated neighbour_index and neighbour_mask.
        """
        if self.neighbour_list_pace == 0:
            pass
        elif steps % self.neighbour_list_pace == 0:
            if self.system_with_energy is not None:
                self.system_with_energy.update_neighbour_list()
            if self.system_with_force is not None:
                self.system_with_force.update_neighbour_list()

            
    def update_bias(self, step: int):
        r"""update bias potential.

        Args:
            step(int):  current step number.
        """
        if self.system_with_energy is not None:
            self.system_with_energy.update_bias(step)

    def update_wrapper(self, step: int):
        r"""update wrapper.

        Args:
            step(int):  current step number.
        """
        if self.system_with_energy is not None:
            self.system_with_energy.update_wrapper(step)

    def update_modifier(self, step: int):
        r"""update modifier.

        Args:
            step(int):  current step number.
        """
        if self.system_with_force is not None:
            self.system_with_force.update_modifier(step)

    def set_pbc_grad(self, value: bool):
        r"""
        set if the pbc_box is trainable.

        Args:
            value(bool):  Whether the pbc_box should be updated.
        """
        if self.system_with_energy is not None:
            self.system_with_energy.set_pbc_grad(value)
        if self.system_with_force is not None:
            self.system_with_force.set_pbc_grad(value)

    def closure(self, *inputs):
        
        energy = 0
        if self.system_with_energy is not None:
            energy = self.system_with_energy(*inputs)
            energy.backward(torch.ones_like(energy))

        if self.system_with_force is not None:
            energy = self.system_with_force(energy)
        
        return energy
    
    def forward(self, *inputs) -> Tuple[Tensor, Tensor]:
        """calculate the energy and the force based on the current coordinates,
        and optimize (update) the coordinates to the next step.

        Returns:
            - energy, Tensor of shape `(B, 1)`. Data type is float.
                     Total potential energy.
            - force, Tensor of shape `(B, A, D)`. Data type is float.
                    Atomic force.

        Symbols:
            B:  Batchsize, i.e., number of walkers of the simulation.
            A:  Number of atoms in the simulation system.
            D:  Spatial dimension of the simulation system. Usually is 3.
        """

        # if self.system_with_energy is not None:
        #     energy = self.system_with_energy(*inputs)
        #     energy.backward(torch.ones_like(energy),retain_graph=True)

        # if self.system_with_force is not None:
        #     energy = self.system_with_force(energy, force, virial)

        # 应用优化器
        if self.optimizer is not None:
            self.optimizer.zero_grad(set_to_none=False)
            self.optimizer.step(lambda: self.closure(*inputs))
        return True
